package limiter

import (
	"context"
	"fmt"
	"time"

	"github.com/go-redis/redis/v8"
)

// RedisRateLimiter 基于Redis的分布式限流器
type RedisRateLimiter struct {
	client    *redis.Client
	keyPrefix string
	limit     int
	period    time.Duration
	luaScript *redis.Script
}

// NewRedisRateLimiter 创建Redis限流器
func NewRedisRateLimiter(client *redis.Client, keyPrefix string, limit int, period time.Duration) *RedisRateLimiter {
	return &RedisRateLimiter{
		client:    client,
		keyPrefix: keyPrefix,
		limit:     limit,
		period:    period,
		luaScript: redis.NewScript(`
		local key = KEYS[1]
		local limit = tonumber(ARGV[1])
		local period = tonumber(ARGV[2])
		local now = tonumber(ARGV[3])
		
		-- 特殊情况处理
		if limit <= 0 then
			return 0 -- 直接拒绝请求
		end
		
		-- 获取当前计数器
		local count = redis.call('GET', key)
		if count == false then
			-- 如果计数器不存在，创建新计数器
			redis.call('SET', key, 1, 'EX', period)
			return 1
		end
		
		-- 计数器已存在，判断是否超过限制
		count = tonumber(count)
		if count < limit then
			-- 未超过限制，增加计数器
			redis.call('INCR', key)
			return count + 1
		end
		
		-- 已超过限制
		return 0
	`),
	}
}

// Allow 检查请求是否允许通过
func (r *RedisRateLimiter) Allow() bool {
	return r.AllowWithKey("")
}

// AllowWithKey 使用特定的键检查请求是否允许通过
func (r *RedisRateLimiter) AllowWithKey(key string) bool {
	// 处理边界情况
	if r.limit <= 0 {
		return false // 直接拒绝请求
	}

	// 检查Redis客户端是否可用
	if r.client == nil {
		return false // 没有客户端时拒绝请求
	}

	if key == "" {
		key = "global"
	}
	fullKey := fmt.Sprintf("%s:%s", r.keyPrefix, key)
	now := time.Now().Unix()

	// 执行Lua脚本
	ctx := context.Background()

	// 尝试ping一下Redis，确认连接正常
	_, pingErr := r.client.Ping(ctx).Result()
	if pingErr != nil {
		return false // Redis连接不可用时拒绝请求
	}

	// 设置有效的period值（至少1秒）
	periodSec := int(r.period.Seconds())
	if periodSec <= 0 {
		// 当period为0时，特殊处理：允许指定数量的请求通过
		// 使用固定key保存计数，不设置过期时间
		count, err := r.client.Incr(ctx, fullKey).Result()
		if err != nil {
			return false
		}
		// 允许前limit个请求通过
		return count <= int64(r.limit)
	}

	result, err := r.luaScript.Run(ctx, r.client, []string{fullKey}, r.limit, periodSec, now).Int()
	if err != nil {
		// 发生错误时，拒绝请求
		return false
	}

	// 结果大于0表示请求允许通过
	return result > 0
}

// UserRateLimiter 用户级别限流器
type UserRateLimiter struct {
	baseLimiter *RedisRateLimiter
}

// NewUserRateLimiter 创建用户级别限流器
func NewUserRateLimiter(client *redis.Client, limit int, period time.Duration) *UserRateLimiter {
	return &UserRateLimiter{
		baseLimiter: NewRedisRateLimiter(client, "ratelimit:user", limit, period),
	}
}

// AllowUser 检查特定用户ID的请求是否允许通过
func (u *UserRateLimiter) AllowUser(userID string) bool {
	return u.baseLimiter.AllowWithKey(userID)
}

// IPRateLimiter IP级别限流器
type IPRateLimiter struct {
	baseLimiter *RedisRateLimiter
}

// NewIPRateLimiter 创建IP级别限流器
func NewIPRateLimiter(client *redis.Client, limit int, period time.Duration) *IPRateLimiter {
	return &IPRateLimiter{
		baseLimiter: NewRedisRateLimiter(client, "ratelimit:ip", limit, period),
	}
}

// AllowIP 检查特定IP的请求是否允许通过
func (i *IPRateLimiter) AllowIP(ipAddress string) bool {
	return i.baseLimiter.AllowWithKey(ipAddress)
}
